# 该脚本用于从图片中提取固定位置的40个像素点并输出为RGB565格式的数组
import sys,os,re
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

# 定义黑色判断阈值 (0-255)
BLACK_THRESHOLD = 50  # 当RGB所有分量都小于此值，判定为纯黑
BLACK_COLOR = 0x0000  # RGB565的纯黑值

def is_black(r, g, b, threshold=BLACK_THRESHOLD):
    """判断颜色是否足够黑"""
    return r < threshold and g < threshold and b < threshold

def rgb888_to_rgb565(r, g, b):
    """RGB888到RGB565转换函数"""
    r = int(r)
    g = int(g)
    b = int(b)
    r = max(0, min(255, r))
    g = max(0, min(255, g))
    b = max(0, min(255, b))
    r5 = (r >> 3) & 0x1F  # 8位红色分量转换为5位
    g6 = (g >> 2) & 0x3F  # 8位绿色分量转换为6位
    b5 = (b >> 3) & 0x1F  # 8位蓝色分量转换为5位
    return (r5 << 11) | (g6 << 5) | b5  # 组合成RGB565格式

def process_texture_image(img, mask):
    """
    处理材质图片并生成RGB565数组和预览图
    
    参数:
        image_path (str): 材质图片路径
        mask (list): 二维掩码矩阵,1表示有效像素
    
    返回:
        rgb565_data (list): RGB565格式的有效像素数据列表
        preview_image (PIL.Image): 预览图像对象
    """
    # 1. 加载图片
    img_width, img_height = img.size
    
    # 2. 检查图片尺寸是否符合要求
    if img_width != 160 or img_height != 160:
        raise ValueError(f"图片尺寸必须为160x160像素,但当前是{img_width}x{img_height}像素")
    
    # 3. 将掩码转换为NumPy数组
    mask_array = np.array(mask, dtype=bool)
    mask_height, mask_width = mask_array.shape
    
    # 4. 计算材质像素块大小 (N x N)
    n_size = img_width // mask_width
    if img_width % mask_width != 0 or img_height % mask_height != 0:
        raise ValueError("图片尺寸必须能被掩码尺寸整除")
    
    # 5. 提取有效像素并转换为RGB565格式
    pixels = np.array(img)
    rgb565_data = []
    
    # 创建预览图像（原始大小）
    preview_pixels = np.zeros_like(pixels)
    
    # 遍历每个材质像素块
    for y in range(mask_height):
        for x in range(mask_width):
            if mask_array[y, x]:
                # 获取当前块的像素区域
                block_x = x * n_size
                block_y = y * n_size
                
                # 取块中心的颜色作为代表颜色
                center_x = block_x + n_size // 2
                center_y = block_y + n_size // 2
                
                # 获取RGB值
                r, g, b = pixels[center_y, center_x]

                # 应用黑色判断规则
                if is_black(r, g, b):
                    rgb565 = BLACK_COLOR
                else:
                    # 转换为RGB565格式
                    rgb565 = rgb888_to_rgb565(r,g,b)

                rgb565_data.append(rgb565)
                
                # 在预览图像中标记有效块（使用原始颜色）
                preview_pixels[block_y:block_y+n_size, block_x:block_x+n_size] = pixels[block_y:block_y+n_size, block_x:block_x+n_size]
            else:
                # 在预览图像中标记无效块为绿色（可选）
                block_x = x * n_size
                block_y = y * n_size
                preview_pixels[block_y:block_y+n_size, block_x:block_x+n_size] = [0, 255, 0]
    
    # 6. 创建预览图像
    preview_img = Image.fromarray(preview_pixels, 'RGB')
    
    return rgb565_data, preview_img

def generate_c_array_single(rgb565_data):
    """生成C语言格式的数组字符串"""
    hex_values = [f"0x{val:04X}" for val in rgb565_data]
    array_str = "const uint16_t texture_data[] = {\n    " + ",\n    ".join(
        ", ".join(hex_values[i:i+8]) for i in range(0, len(hex_values), 8)
    ) + "\n};\n"
    
    array_info = f"// 数组大小: {len(rgb565_data)} 像素\n"
    return array_info + array_str

def generate_c_array(all_texture_data):
    """生成C语言二维数组格式的字符串"""
    # 计算图像数量和每张图像的有效像素数
    num_images = len(all_texture_data)
    if num_images == 0:
        return ""
    
    pixels_per_image = len(all_texture_data[0])
    
    # 检查所有图像是否具有相同数量的有效像素
    for i, data in enumerate(all_texture_data):
        if len(data) != pixels_per_image:
            raise ValueError(f"图像 {i} 的有效像素数({len(data)})与其他图像不同({pixels_per_image})")
    
    # 生成C数组头
    result = "const uint16_t texture_data[][{pixels_per_image}] = {{\n"
    
    # 添加每张图像的数组数据
    for i, data in enumerate(all_texture_data):
        hex_values = [f"0x{val:04X}" for val in data]
        line = "    { " + ", ".join(hex_values) + " }"
        if i < num_images - 1:
            line += ",  // Image " + str(i)
        else:
            line += "   // Image " + str(i)
        result += line + "\n"
    
    result += "};\n"
    
    # 添加注释信息
    result = f"// 共 {num_images} 张图像，每张 {pixels_per_image} 个有效像素\n" + result
    result += f"#define NUM_TEXTURE_IMAGES {num_images}\n"
    result += f"#define PIXELS_PER_IMAGE {pixels_per_image}\n"
    
    return result

def process_single_image(image_path,mask):
    """处理一张图片"""
    # 处理图片
    img = Image.open(image_path).convert('RGB')
    rgb565_data, preview_img = process_texture_image(img, mask)
    
    # 显示预览图像
    plt.figure(figsize=(8, 8))
    plt.imshow(preview_img)
    plt.title("valid area(invalid red)")
    plt.axis('off')
    plt.show()
    
    # 打印提取的有效像素数量
    print(f"提取的有效像素数量: {len(rgb565_data)}")
    
    # 生成并打印C数组
    c_array_str = generate_c_array_single(rgb565_data)
    print("\nC语言数组:\n")
    print(c_array_str)
    
    # 保存C数组到文件
    with open("texture_data.h", "w") as f:
        f.write(c_array_str)
    
    # 保存预览图像
    preview_img.save("texture_preview.png")
    print("预览图像已保存为 %s'texture_preview.png'"%(os.getcwd()))

def process_all_images(input_dir,mask):
    """批量处理文件夹中的所有Compass_*.png图像"""
    # 获取匹配的文件列表并排序
    all_files = [f for f in os.listdir(input_dir) if f.startswith('Compass_') and f.endswith('.png')]
    all_files.sort(key=lambda x: int(re.search(r'Compass_(\d+)\.png', x).group(1)))
    
    if not all_files:
        print("警告:没有找到任何Compass_*.png文件")
        return [], []
    
    # 为每张图像创建输出文件夹
    output_dir = os.path.join(input_dir, "output")
    preview_dir = os.path.join(output_dir, "previews")
    os.makedirs(preview_dir, exist_ok=True)
    
    all_texture_data = []
    preview_images = []
    
    # 处理每张图像
    for file in all_files:
        file_path = os.path.join(input_dir, file)
        try:
            # 加载图像
            img = Image.open(file_path).convert('RGB')
            
            # 处理图像
            rgb565_data, preview_img = process_texture_image(img, mask)
            
            # 保存结果
            all_texture_data.append(rgb565_data)
            preview_images.append(preview_img)
            
            # 保存预览图像
            preview_file = f"preview_{file}"
            preview_img.save(os.path.join(preview_dir, preview_file))
            
            print(f"已处理: {file} -> {len(rgb565_data)}个有效像素")
        except Exception as e:
            print(f"处理 {file} 时出错: {str(e)}")
    
    #return all_texture_data, preview_images
    if all_texture_data:
        # 生成并保存C语言数组
        c_array_str = generate_c_array(all_texture_data)
        output_path = os.path.join(input_dir, "output", "texture_data.h")
        with open(output_path, "w") as f:
            f.write(c_array_str)
        
        print(f"\nC语言数组已保存至: {output_path}")
        
        # 显示第一张预览图
        plt.figure(figsize=(8, 8))
        plt.imshow(preview_images[0])
        plt.title("valid area(invalid red)")
        plt.axis('off')
        plt.show()
    else:
        print("没有生成数据")

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("用法:python extract_pic.py <输入PIC文件(夹)>")
        sys.exit(1)

    image_path = sys.argv[1]

    mask = [ # 指针区域
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], # start
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
        [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
        [0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
        [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], # end
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    ]
    mask_inverted = [ # 其他区域
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0], # start
    [0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
    [0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
    [0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0],  
    [0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0],  
    [0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0],  
    [0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0],  
    [0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 0],  
    [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0],
    [0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
    [0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
    [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0], # end
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
]

    if os.path.isfile(image_path):
        process_single_image(image_path,mask) # 处理单张图片-输入文件
    elif os.path.isdir(image_path):
        process_all_images(image_path,mask) # 批量处理-输入文件夹
    
    

    
